from label_studio_ml.model import LabelStudioMLBase
import os
import logging
import numpy as np
import cv2
import requests
from ultralytics import YOLO

logger = logging.getLogger(__name__)

class YOLOv8Model(LabelStudioMLBase):
    def __init__(self, **kwargs):
        super(YOLOv8Model, self).__init__(**kwargs)
        
        # Get configuration from kwargs
        self.from_name = kwargs.get('from_name', 'image')
        self.to_name = kwargs.get('to_name', 'label')
        self.labels = self.parsed_label_config.get('label').get('labels')
        self.model_path = 'models/yolo11n_dnf.pt'
        self.model_version = self.model_path
        
        try:
            # Load YOLO model, use yolov8n.pt as default model
            
            if not os.path.exists(self.model_path):
                logger.warning(f'Model file {self.model_path} not found, will download default model')
            
            self.model = self._load_model()
            logger.info(f'Model loaded successfully from {self.model_path}')
        except Exception as e:
            logger.error(f'Error loading model: {str(e)}')
            self.model = None
        
    def _load_model(self):
        model = YOLO(self.model_path)
        return model
    
    def validate(self, tasks, context=None, **kwargs):
        """Validate task data and check if model can make predictions"""
        if self.model is None:
            return False
        return True
        
    def predict(self, task, **kwargs):
        if self.model is None:
            logger.error('Model is not properly initialized')
            return []
            
        regions  = []
        try:
            # Get image path or URL
            image_path = task['data'].get(self.from_name)
            # Read image
            image = self._get_image(image_path)
            # Run prediction
            results = self.model(image)
            # Get class names
            model_names = self.model.names
            # Get predictions
            data = results[0].boxes
            for i in range(data.shape[0]):  # iterate over items
                score = float(data.conf[i])  # tensor => float
                x, y, w, h = data.xywhn[i].tolist()
                model_label = model_names[int(data.cls[i])]
                # output_label = self.labels[model_label]
                region = {
                    "from_name": 'label',
                    "to_name": 'image',
                    "type": "rectanglelabels",
                    "value": {
                        "rectanglelabels": [model_label],
                        "x": (x - w / 2) * 100,
                        "y": (y - h / 2) * 100,
                        "width": w * 100,
                        "height": h * 100,
                    },
                    "score": score,
                }
                regions.append(region)
        except Exception as e:
            logger.error(f'Error processing task: {str(e)}')

        logger.info(f'Predictions: {regions}')
        return regions
    
    def _get_image(self, image_path):
        try:
            response= requests.get('http://localhost:8080'+image_path, headers={'Authorization': 'Token 8f470645906dd75b6ea3e96d23eaac6880a42373'})
            print(response)
            image_array = np.asarray(bytearray(response.content), dtype=np.uint8)
            image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
            return image
        except Exception as e:
            logger.error(f'Error reading image {image_path}: {str(e)}')
            return None
